import asyncio
import websockets
import numpy as np
import threading
import socket
import time
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import CompressedImage
from ai_msgs.msg import PerceptionTargets
from cv_bridge import CvBridge
import cv2
import traceback
from message_filters import Subscriber, ApproximateTimeSynchronizer
# 全局变量存储最新图像和统计信息
latest_frame = None
frame_lock = threading.Lock()
stats = {}
stats_lock = threading.Lock()

# 标签映射表（示例）
LABEL_MAP = {
    "person": "人",
    "car": "车",
    "dog": "狗",
    # 添加更多标签映射
}

# 动态获取本机 IP 地址
def get_local_ip():
    try:
        # 创建一个 UDP 套接字连接到公共 DNS 服务器（如 Google DNS）
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.connect(("8.8.8.8", 80))  # 连接到外部地址以获取本地 IP
            local_ip = s.getsockname()[0]  # 获取本地 IP 地址
        return local_ip
    except Exception as e:
        print(f"无法获取本地 IP 地址: {e}")
        return "127.0.0.1"  # 如果失败，返回默认本地地址

# ROS2 图像订阅节点
class ImageSubscriber(Node):
    def __init__(self):
        super().__init__('image_subscriber')
        self.bridge = CvBridge()
        self.callback_counter = 0  # 初始化计数器

        # 订阅图像和检测结果话题
        self.image_sub = Subscriber(self, CompressedImage, '/image_raw/compressed')
        self.detection_sub = Subscriber(self, PerceptionTargets, '/hobot_dnn_detection')

        # 时间同步
        self.ts = ApproximateTimeSynchronizer(
            [self.image_sub, self.detection_sub],
            queue_size=10,
            slop=0.5  # 时间容忍范围
        )
        self.ts.registerCallback(self.callback)

        # 心跳机制
        self.timer = self.create_timer(5.0, self.check_health)  # 每 5 秒检查一次健康状态

        # 健康状态标志
        self.last_callback_time = time.time()  # 上次回调时间
        self.get_logger().info('this Web Node started.')

    def callback(self, img_msg, det_msg):
        try:
            # 更新健康状态
            self.last_callback_time = time.time()


            np_arr = np.frombuffer(img_msg.data, np.uint8)
            cv_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)

            # 更新全局变量 latest_frame
            with frame_lock:
                global latest_frame
                latest_frame = cv_image
                current_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
                self.get_logger().info(f"Callback: Latest frame updated at [{current_time}].")
        except Exception as e:
            # 使用 error 级别，并打印详细 Traceback
            self.get_logger().error(f"Exception during callback processing: {e}")


    def check_health(self):
        current_time = time.time()
        time_since_last_callback = current_time - self.last_callback_time

        if time_since_last_callback > 10:  # 如果超过 10 秒未触发回调
            try:
                self.resubscribe_topics()
                self.get_logger().info("Health check: Resubscription attempt finished.")
            except Exception as e:
                self.get_logger().error(f"Health check: Exception during resubscription attempt: {e}")
        else:
            self.get_logger().info("Health check: Node is still active (callbacks within 10s).")

    def resubscribe_topics(self):
        try:
            self.get_logger().info("Resubscribe: Attempting to unregister existing subscribers...")
            try:
                self.image_sub.unregister()
                self.get_logger().info("Resubscribe: image_sub unregistered.")
            except Exception as e:
                self.get_logger().warn(f"Resubscribe: Failed to unregister image_sub (might be ok if already invalid): {e}")
            try:
                self.detection_sub.unregister()
                self.get_logger().info("Resubscribe: detection_sub unregistered.")
            except Exception as e:
                self.get_logger().warn(f"Resubscribe: Failed to unregister detection_sub (might be ok if already invalid): {e}")

            self.image_sub = Subscriber(self, CompressedImage, '/image_raw/compressed')
            self.detection_sub = Subscriber(self, PerceptionTargets, '/hobot_dnn_detection')

            self.ts = ApproximateTimeSynchronizer(
                [self.image_sub, self.detection_sub],
                queue_size=10,
                slop=0.5
            )
            self.ts.registerCallback(self.callback)
            self.get_logger().info("Topics resubscribed successfully!")
        except Exception as e:
            self.get_logger().error(f"Failed during resubscription process: {e}")
            self.get_logger().error(traceback.format_exc())

# WebSocket 视频流发送函数
async def send_video(websocket):
    global latest_frame
    try:
        while True:
            frame_to_send = None
            with frame_lock:
                if latest_frame is not None:
                    frame_to_send = latest_frame.copy()  # 复制以快速释放锁
                    memory_size = frame_to_send.nbytes  # 获取图像数据的字节大小
            if frame_to_send is not None:
                # 编码为 JPEG
                ret, jpeg = cv2.imencode('.jpg', frame_to_send, [cv2.IMWRITE_JPEG_QUALITY, 80])
                if ret:
                    await websocket.send(jpeg.tobytes())  # 发送 JPEG 数据

            # 控制帧率 (20 FPS)
            await asyncio.sleep(1 / 20)
    except websockets.exceptions.ConnectionClosed:
        print("WebSocket client disconnected.")
    except Exception as e:
        print(f"Exception in WebSocket send_video loop: {e}")
        traceback.print_exc()

# 主函数
async def main():
    print("Main: Starting rclpy.init()...")
    rclpy.init()
    print("Main: rclpy.init() finished.")

    try:
        print("Main: Creating ImageSubscriber node...")
        yolo_web_node = ImageSubscriber()
        print("Main: ImageSubscriber node created.")

        print("Main: Starting ROS 2 event loop thread...")
        ros_thread = threading.Thread(target=rclpy.spin, args=(yolo_web_node,), daemon=True)
        ros_thread.start()
        print(f"Main: ROS 2 event loop thread started with name: {ros_thread.name}, daemon: {ros_thread.daemon}")

        print("Main: Getting local IP address...")
        local_ip = get_local_ip()
        print(f"Main: Local IP address: {local_ip}")
        print(f"WebSocket 服务器将在 {local_ip}:3001 上运行")

        print("Main: Starting WebSocket server...")
        async with websockets.serve(send_video, '0.0.0.0', 3001):
            print("WebSocket 服务器已启动")
            await asyncio.Future()  # 持续运行 WebSocket 服务器
        print("Main: asyncio Future wait finished (this should not happen in normal run).")
    finally:
        if 'yolo_web_node' in locals() and yolo_web_node is not None:
            yolo_web_node.destroy_node()
            print("Main: ROS 2 node destroyed.")
        print("Main: Shutting down rclpy...")
        rclpy.shutdown()


if __name__ == '__main__':
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("程序已退出 (KeyboardInterrupt)")
    except Exception as e:
        print(f"An unhandled exception occurred in main: {e}")
        traceback.print_exc()
